import pylab as P
import pygtk
import gtk
import gobject

def print_fig_prop(fig):
    for ind_ax, ax in enumerate(P.figure(fig).axes):
        print "*** Axe %d ***"%ind_ax
        print "Xmin : %g"%ax.get_xlim()[0]
        print "Xmax : %g"%ax.get_xlim()[1]
        print "Ymin : %g"%ax.get_ylim()[0]
        print "Ymax : %g"%ax.get_ylim()[1]
        print "Fontsize (X) : %g"%ax.get_xticklabels()[0].get_fontsize()
        print "Fontsize (Y) : %g"%ax.get_yticklabels()[0].get_fontsize()
        print "Title : " + ax.title.get_text()
        print "Title fontsize : %g"%ax.title.get_fontsize()
        print "X label : " + ax.xaxis.get_label().get_text()
        print "X label fontsize : %g"%ax.xaxis.get_label().get_fontsize()
        print "Y label : " + ax.yaxis.get_label().get_text()
        print "Y label fontsize : %g"%ax.yaxis.get_label().get_fontsize()
        
        for ind_line, line in enumerate(ax.lines):
            print "  *** Line %d ***"%ind_line
            print "  Style : " + line.get_linestyle()
            print "  Marker : " + line.get_marker()
	    print "  Line Width : %.1f"%line.get_linewidth()
	    print "  Color : " + line.get_color()
        for ind_im, im in enumerate(ax.images):
            print "  *** Image %d ***"%ind_im
            print "  Colormap : " + im.cmap.name
            print "  Interpolation : " + im.get_interpolation()
            print "  Extent : [%g, %g, %g, %g]"%(im.get_extent()[0],im.get_extent()[1],im.get_extent()[2],im.get_extent()[3], )
        
class properties:

    def row_activated(self, treeview):
        model, iter = treeview.get_selection().get_selected()
        type = model.get_value(iter, 1)
        self.type_selected = type
        obj = model.get_value(iter,2)
        self.selected = model.get_value(iter,2)
        self.fig = model.get_value(iter, 3)
        if type == 'line':
            self.display_line(obj)
        elif type == 'axe':
            self.display_axe(obj)
        elif type == 'image':
            self.display_image(obj)
        elif type == 'figure':
            self.display_figure(obj)
        elif type == 'text':
            self.display_text(obj)
##         if type != 'figure':
##             visi = obj.get_visible()
##             obj.set_visible(not(visi))
##             P.figure(self.fig)
##             P.draw()
##             obj.set_visible(visi)
##             P.figure(self.fig)
##             P.draw()

    def get_value(self, obj, name):
        if name == 'linestyle':
            return obj.get_linestyle()
        elif name == 'marker':
            return obj.get_marker()
        elif name == 'linewidth':
            return obj.get_linewidth()
        elif name == 'color':
            return obj.get_color()
        elif name == 'xmin':
            return obj.get_xlim()[0]
        elif name == 'xmax':
            return obj.get_xlim()[1]
        elif name == 'ymin':
            return obj.get_ylim()[0]
        elif name == 'ymax':
            return obj.get_ylim()[1]
        elif name == 'title':
            return obj.title.get_text()
        elif name == 'titlefontsize':
            return obj.title.get_fontsize()
        elif name == 'xlabel':
            return obj.xaxis.get_label().get_text()
        elif name == 'ylabel':
            return obj.yaxis.get_label().get_text()
        elif name == 'xfontsize':
            return obj.xaxis.get_ticklabels()[0].get_fontsize()
        elif name == 'yfontsize':
            return obj.yaxis.get_ticklabels()[0].get_fontsize()
        elif name == 'xlabelfontsize':
            return obj.xaxis.get_label().get_fontsize()
        elif name == 'ylabelfontsize':
            return obj.yaxis.get_label().get_fontsize()
        elif name == 'figwidth':
##            return obj.window.get_size()[0]
            return obj.get_figwidth()
        elif name == 'figheight':
##            return obj.window.get_size()[1]
            return obj.get_figheight()
        elif name == 'xpos':
            return obj.get_position()[0]
        elif name == 'ypos':
            return obj.get_position()[1]
        elif name == 'width':
            return obj.get_position()[2]
        elif name == 'height':
            return obj.get_position()[3]
        elif name == 'xbegin':
            return obj.get_extent()[0]
        elif name == 'xend':
            return obj.get_extent()[1]
        elif name == 'ybegin':
            return obj.get_extent()[2]
        elif name == 'yend':
            return obj.get_extent()[3]
        elif name == 'cmap':
            return obj.cmap.name
        elif name == 'interp':
            return obj.get_interpolation()
        elif name == 'label':
            return obj.get_label()
        elif name == 'textx':
            return obj.get_position()[0]
        elif name == 'texty':
            return obj.get_position()[1]
        elif name == 'text':
            return obj.get_text()
        elif name == 'textfontsize':
            return obj.get_fontsize()
        elif name == 'textrotation':
            return obj.get_rotation()
        elif name == 'markersize':
            return obj.get_markersize()
        elif name == 'markerfacecolor':
            return obj.get_markerfacecolor()
        elif name == 'markeredgecolor':
            return obj.get_markeredgecolor()
        elif name == 'markeredgewidth':
            return obj.get_markeredgewidth()

        
    def set_value(self, obj, name, value):
        if name == 'linestyle':
            obj.set_linestyle(value)
        elif name == 'marker':
            obj.set_marker(value)
        elif name == 'linewidth':
            obj.set_linewidth(float(value))
        elif name == 'color':
            obj.set_color(value)
        elif name == 'xmin' or name == 'xmax':
            obj.set_xlim([float(self.param['xmin'].get_text()), float(self.param['xmax'].get_text())])
        elif name == 'ymin' or name == 'ymax':
            obj.set_ylim([float(self.param['ymin'].get_text()), float(self.param['ymax'].get_text())])
        elif name == 'title':
            obj.title.set_text(value)
        elif name == 'titlefontsize':
            obj.title.set_fontsize(float(value))
        elif name == 'xlabel':
            obj.xaxis.get_label().set_text(value)
        elif name == 'ylabel':
            obj.yaxis.get_label().set_text(value)
        elif name == 'xfontsize':
            for tick in obj.xaxis.get_ticklabels():
                tick.set_fontsize(float(value))
        elif name == 'yfontsize':
            for tick in obj.yaxis.get_ticklabels():
                tick.set_fontsize(float(value))
        elif name == 'xlabelfontsize':
            obj.xaxis.get_label().set_fontsize(float(value))
        elif name == 'ylabelfontsize':
            obj.yaxis.get_label().set_fontsize(float(value))
##        elif name == 'figwidth' or name == 'figheight':
##            obj.resize(int(self.param['figwidth'].get_text()),int(self.param['figheight'].get_text()) )
        elif name == 'figwidth':
            obj.set_figwidth(float(value))
        elif name == 'figheight':
            obj.set_figheight(float(value))
        elif name == 'xpos' or name =='ypos' or name == 'width' or name == 'height':
            a = float(self.param['xpos'].get_text())
            b = float(self.param['ypos'].get_text())
            c = float(self.param['width'].get_text())
            d = float(self.param['height'].get_text())
            obj.set_position([a,b,c,d])
        elif name == 'cmap':
            obj.set_cmap(getattr(P.cm, value))
        elif name == 'interp':
            obj.set_interpolation(value)
        elif name == 'label':
            obj.set_label(value)
        elif name == 'textx':
            obj.set_x(float(value))
        elif name == 'texty':
            obj.set_y(float(value))
        elif name == 'text':
            obj.set_text(value)
        elif name == 'textfontsize':
            obj.set_fontsize(float(value))
        elif name == 'textrotation':
            obj.set_rotation(float(value))
        elif name == 'markersize':
            obj.set_markersize(float(value))
        elif name == 'markerfacecolor':
            obj.set_markerfacecolor(value)
        elif name == 'markeredgecolor':
            obj.set_markeredgecolor(value)
        elif name == 'markeredgewidth':
            obj.set_markeredgewidth(float(value))

    def remove(self, widget):
        self.paramtable.remove(widget)

    def add_row(self, name, title, obj, line):
        tmp = gtk.Label(title)
        tmp.show()
        self.paramtable.attach(tmp,0,1,line,line+1, yoptions=gtk.FILL)

        default = self.get_value(obj, name)
##         if default == None or default == 'None':
##             default = ''
        if not self.param_treestore.has_key(name):
            self.param[name] = gtk.Entry(255)
            self.param[name].set_text(str(default))
            self.param[name].connect("activate", self.apply)
        else:
            self.param[name] = gtk.ComboBox(self.param_treestore[name])
            cell = gtk.CellRendererText()
            self.param[name].pack_start(cell, True)
            self.param[name].add_attribute(cell, 'text', 0)
            ii=0
            while self.param_treestore[name][ii][0] != default:
                ii = ii+1
            self.param[name].set_active(ii)
            self.param[name].connect("changed", self.apply)
        self.param[name].show()
        self.paramtable.attach(self.param[name], 1, 2, line, line+1, yoptions=gtk.FILL)

    def display_line(self, line):
        self.paramtable.foreach(self.remove)
        self.param.clear()
        self.paramtable.resize(2, len(self.param_lines_order))
        for ii, name in enumerate(self.param_lines_order):
            self.add_row(name, self.param_names[name], line, ii)
            
    def display_axe(self, ax):
        self.paramtable.foreach(self.remove)
        self.param.clear()
        self.paramtable.resize(2, len(self.param_axe_order))
        for ii,name in enumerate(self.param_axe_order):
            self.add_row(name, self.param_names[name], ax, ii)

    def display_figure(self, figure):
        self.paramtable.foreach(self.remove)
        self.param.clear()
        self.paramtable.resize(2, len(self.param_figure_order))
        for ii, name in enumerate(self.param_figure_order):
            self.add_row(name, self.param_names[name], figure, ii)

    def display_text(self, text):
        self.paramtable.foreach(self.remove)
        self.param.clear()
        self.paramtable.resize(2, len(self.param_text_order))
        for ii, name in enumerate(self.param_text_order):
            self.add_row(name, self.param_names[name], text, ii)

    def display_image(self, im):
        self.paramtable.foreach(self.remove)
        self.param.clear()
        self.paramtable.resize(2, len(self.param_image_order))
        for ii, name in enumerate(self.param_image_order):
            self.add_row(name, self.param_names[name], im, ii)

    def apply(self, widget):
        obj = self.selected
        for name, button in self.param.items():
            if not self.param_treestore.has_key(name):
                value = button.get_text()
            else:
                iter = button.get_active_iter()
                value = self.param_treestore[name].get_value(iter, 0)
            self.set_value(obj, name, value)
        P.figure(self.fig)
        P.draw()
        self.row_activated(self.treeview)

    def save(self, widget):
        model, iter = self.treeview.get_selection().get_selected()
        if self.type_selected == 'line':
            line = model.get_value(iter,2)
            x = P.transpose(P.array([line.get_xdata(), line.get_ydata()]))
            fileselection = gtk.FileSelection('Save ' + model.get_value(iter,0) + ' data in...')
            fileselection.ok_button.connect("clicked", self.go_save, fileselection, x)
            fileselection.cancel_button.connect("clicked", lambda w: fileselection.destroy())
            fileselection.show()
    def go_save(self, widget, win, data):
        name = win.get_filename()
        P.save(name, data)
        win.destroy()

    def delete(self, widget):
        model, iter = self.treeview.get_selection().get_selected()
        if self.type_selected=='axe':
            P.delaxes(self.selected)
            P.figure(self.fig)
            P.draw()
        elif self.type_selected=='figure':
            P.close(self.fig)
        elif self.type_selected=='line':
            axe = model.get_value(model.iter_parent(iter), 2)
            del axe.lines[model.get_value(iter, 4)]
            P.figure(self.fig)
            P.draw()
        elif self.type_selected=='image':
            axe = model.get_value(model.iter_parent(iter), 2)
            del axe.images[model.get_value(iter, 4)]
            P.figure(self.fig)
            P.draw()
        elif self.type_selected=='text':
            axe = model.get_value(model.iter_parent(iter), 2)
            del axe.texts[model.get_value(iter, 4)]
            P.figure(self.fig)
            P.draw()
        self.treestore.remove(iter)
        self.paramtable.foreach(self.remove)
        self.param.clear()
        self.update(None)

    def update(self, widget):
        self.treestore.clear()
        for figure in P.matplotlib._pylab_helpers.Gcf.get_all_fig_managers():
            fig = figure.num
            iter_fig = self.treestore.append(None, ['Figure %d'%fig, 'figure' ,P.figure(fig), fig, fig])
            for ind_ax, ax in enumerate(P.figure(fig).axes):
                iter_ax = self.treestore.append(iter_fig, ['Axe %d'%ind_ax, 'axe' ,ax, fig, ind_ax])
                for ind_line, line in enumerate(ax.lines):
                    iter = self.treestore.append(iter_ax, ['Line %d'%ind_line, 'line', line, fig, ind_line])
                for ind_im, im in enumerate(ax.images):
                    iter = self.treestore.append(iter_ax, ['Image %d'%ind_im, 'image',im, fig, ind_im])
                for ind_text, text in enumerate(ax.texts):
                    iter = self.treestore.append(iter_ax, ['Text %d'%ind_text, 'text', text, fig, ind_text])
        try:
            self.fig=fig
        except:
            self.window.destroy()

    def copy(self, widget):
        model, iter = self.treeview.get_selection().get_selected()
        type_to_copy = model.get_value(iter, 1)
        element_to_copy = model.get_value(iter, 2)
        if type_to_copy == 'figure':
            return
        self.wincopy = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.wincopy.set_transient_for(self.window)
        self.wincopy.set_modal(True)
        self.wincopy.set_title('Copy ' + model.get_value(iter, 0) + ' to...')
        self.wincopy.set_default_size(200,100)
        vbox = gtk.VBox()
        self.wincopy.add(vbox)
        scroll = gtk.ScrolledWindow()
        scroll.set_policy(gtk.POLICY_NEVER, gtk.POLICY_AUTOMATIC)
        vbox.pack_start(scroll)
        model_filter = model.filter_new()
        model_filter.set_visible_func(self.visible_model, type_to_copy)
        self.treeview_copy = gtk.TreeView(model_filter)
        renderer = gtk.CellRendererText()
        column = gtk.TreeViewColumn('Object ID', renderer, text=0)
        self.treeview_copy.set_headers_visible(False)
        self.treeview_copy.append_column(column)
        self.treeview_copy.show()
        scroll.add(self.treeview_copy)
        scroll.show()
        hbox = gtk.HBox()
        vbox.pack_start(hbox)
        button = gtk.Button('Cancel')
        button.connect("clicked", self.close_copy)
        button.show()
        hbox.pack_start(button)
        button = gtk.Button('Copy')
        button.connect("clicked", self.do_copy, type_to_copy, element_to_copy)
        button.show()
        hbox.pack_start(button)
        hbox.show()
        vbox.show()
        self.wincopy.show()

    def do_copy(self, widget, type, element):
        model_to, iter_to = self.treeview_copy.get_selection().get_selected()
        dest = model_to.get_value(iter_to, 2)
        if type == 'line':
            line = P.Line2D(element.get_xdata(), element.get_ydata(),
                            linewidth=element.get_linewidth(), linestyle=element.get_linestyle(),
                            color=element.get_color(), marker=element.get_marker())
            dest.add_line(line)            
        elif type == 'image':
            dest.images.append(element)
        elif type == 'text':
            dest.text(element.get_position()[0], element.get_position()[1],
                      element.get_text(), fontsize=element.get_fontsize())
        self.close_copy(None)
        self.update(None)
    
    def close_copy(self, widget):
        self.wincopy.destroy()

    def visible_model(self, model, iter, data):
        type = model.get_value(iter, 1)
        if data == 'axe':
            if type == 'figure':
                return True
            else:
                return False
        else:
            if type == 'figure' or type == 'axe':
                return True
            else:
                return False

    def end_of_program(self, widget, data=None):
        if P.backends.backend[:3]=='GTK':
            self.window.destroy()
        else:
            gtk.main_quit()

    def __init__(self):
        self.selected=None
        self.param={}
        marker_types = ['.', ',', 'o', '^', 'v', '<', '>', 's', '+', 'x',
                        'D', 'd', '1', '2', '3', '4', 'h', 'H',
                        'p', '|', '_', 'None']
        marker_treestore = gtk.TreeStore(str)
        for marker in marker_types:
            marker_treestore.append(None, [marker])
        line_types = ['-', '--', '-.', ':', 'steps', 'None']
        line_treestore = gtk.TreeStore(str)
        [line_treestore.append(None, [ii]) for ii in line_types]
        self.param_names = {'label': 'Label',
                            'linestyle': 'Line style',
                            'marker': 'Marker',
                            'markersize': 'Marker Size',
                            'markerfacecolor': 'Marker Face Color',
                            'markeredgecolor': 'Marker Edge Color',
                            'markeredgewidth': 'Marker Edge Width',
                            'linewidth': 'Line Width',
                            'color': 'Color',
                            'cmap': 'Colormap',
                            'interp': 'Interpolation',
                            'label': 'Label',
                            'xpos': 'X position',
                            'ypos': 'Y position',
                            'width': 'Width',
                            'height': 'Height',
                            'xmin': 'X min',
                            'xmax': 'X max',
                            'ymin': 'Y min',
                            'ymax': 'Y max',
                            'title': 'Title',
                            'titlefontsize': 'Title fontsize',
                            'xlabel': 'X label',
                            'ylabel': 'Y label',
                            'xlabelfontsize': 'Fontsize (X label)',
                            'ylabelfontsize': 'Fontsize (Y label)',
                            'xfontsize': 'Fontsize (X ticks)',
                            'yfontsize': 'Fontsize (Y ticks)',
                            'figwidth': 'Width',
                            'figheight': 'Height',
                            'textx': 'X position',
                            'texty': 'Y position',
                            'text': 'Text',
                            'textfontsize': 'Fontsize',
                            'textrotation': 'Rotation'}

        self.param_lines_order = ['label', 'linestyle', 'color', 'linewidth',
                                  'marker', 'markersize', 'markerfacecolor', 'markeredgecolor', 'markeredgewidth']
        self.param_image_order = ['label', 'cmap', 'interp']
        self.param_axe_order = ['xpos', 'ypos', 'width', 'height',
                                'xmin', 'xmax', 'ymin', 'ymax',
                                'title', 'titlefontsize',
                                'xlabel', 'xlabelfontsize',
                                'ylabel', 'ylabelfontsize',
                                'xfontsize', 'yfontsize']
        self.param_figure_order = ['figwidth', 'figheight']
        self.param_text_order = ['textx', 'texty', 'text', 'textfontsize',
                                 'color', 'textrotation']
        self.param_treestore = {'marker': marker_treestore,
                                'linestyle': line_treestore}
                          
        self.window = gtk.Window(gtk.WINDOW_TOPLEVEL)
        self.window.set_title('Graphics properties')
        self.window.set_default_size(300, 400)
        self.window.connect("delete_event", self.end_of_program)
        self.window.connect("destroy", self.end_of_program)
        vbox = gtk.VBox()
        hbox = gtk.HBox()
        self.paramtable = gtk.Table()
        self.paramtable.show()
        self.treestore = gtk.TreeStore(gobject.TYPE_STRING, str, object, int, int)
        self.treeview = gtk.TreeView(self.treestore)
        self.treeview.set_headers_visible(False)
        renderer = gtk.CellRendererText()
        self.column = gtk.TreeViewColumn('Object ID', renderer, text=0)
        self.treeview.append_column(self.column)

        self.treeview.connect("cursor-changed", self.row_activated)
        scroll = gtk.ScrolledWindow()
        scroll.set_policy(gtk.POLICY_NEVER, gtk.POLICY_AUTOMATIC)
        scroll.add(self.treeview)
        scroll.show()
        self.update(None)
        hbox.pack_start(scroll)

        viewport = gtk.Viewport()
        viewport.add(self.paramtable)
        viewport.show()
        scroll = gtk.ScrolledWindow()
        scroll.set_policy(gtk.POLICY_NEVER, gtk.POLICY_AUTOMATIC)
        scroll.add(viewport)
        scroll.show()
        hbox.pack_start(scroll)

        bottom=gtk.HBox()
        button = gtk.Button('Update')
        button.connect("clicked", self.update)
        button.show()
        bottom.pack_start(button)
        button = gtk.Button('Copy')
        button.connect("clicked", self.copy)
        button.show()
        bottom.pack_start(button)
        button = gtk.Button('Delete')
        button.connect("clicked", self.delete)
        button.show()
        bottom.pack_start(button)
        bottom.show()
        button = gtk.Button('Save data')
        button.connect("clicked", self.save)
        button.show()
        bottom.pack_start(button)
        bottom.show()
        button = gtk.Button('Apply')
        button.connect("clicked", self.apply)
        button.show()
        bottom.pack_start(button)
        bottom.show()
        
        vbox.pack_start(hbox)
        vbox.pack_start(bottom, False)
        self.window.add(vbox)
        self.treeview.show()
        hbox.show()
        vbox.show()
        self.window.show()
        gtk.main()
        
